import sys
import struct
from typing import Optional

import numpy as np
import threading
import time
from scipy.io import savemat
from crcmod import mkCrcFun

from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QComboBox, QPushButton, QLineEdit, QLabel, QGroupBox, QStackedWidget)
from PyQt5.QtCore import QObject, Qt, pyqtSignal, pyqtSlot, QThread, QTimer, QMutex
from PyQt5.QtSerialPort import QSerialPort, QSerialPortInfo
import socket
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure


# ================== 通信模块 ==================
class BaseCommunicator(QObject):
    data_received = pyqtSignal(bytes)
    connected = pyqtSignal()
    disconnected = pyqtSignal()
    sn_received = pyqtSignal(str)
    error = pyqtSignal(str)

    def connect(self): pass

    def disconnect(self): pass

    def send_command(self, cmd: str): pass


class SerialCommunicator(BaseCommunicator):
    def __init__(self, port, baudrate):
        super().__init__()
        self.serial = QSerialPort()
        self.serial.setPortName(port)
        self.serial.setBaudRate(baudrate)
        self.serial.readyRead.connect(self._handle_read)
        self.serial.setDataBits(QSerialPort.Data8)
        self.serial.setParity(QSerialPort.NoParity)
        self.serial.setStopBits(QSerialPort.OneStop)

    def connect(self):
        if self.serial.isOpen():
            self.error.emit("串口已打开")
            return
        if self.serial.open(QSerialPort.ReadWrite):
            self.connected.emit()
        else:
            self.error.emit(f"串口打开失败: {self.serial.errorString()}")

    def send_command(self, cmd, terminator="\r\n", encoding="utf-8"):
        data = f"{cmd}{terminator}".encode(encoding)
        if self.serial.isOpen():
            self.serial.write(data)

    def _handle_read(self):
        if self.serial.bytesAvailable():
            data = self.serial.readAll().data()
            self.data_received.emit(data)

    def disconnect(self):
        if self.serial.isOpen():
            self.serial.close()
        self.disconnected.emit()


class WiFiCommunicator(BaseCommunicator):
    def __init__(self, ip: str, port: int):
        super().__init__()
        self.ip = ip
        self.port = port
        self.sock: Optional[socket.socket] = None
        self.connector_thread: Optional[WiFiConnectorThread] = None
        self.listener_thread: Optional[WiFiListenerThread] = None
        self.mutex = QMutex()

    def connect(self):
        self.mutex.lock()
        try:
            if self.connector_thread and self.connector_thread.isRunning():
                return

            self.connector_thread = WiFiConnectorThread(self.ip, self.port)
            self.connector_thread.connected.connect(self._on_connected)
            self.connector_thread.error.connect(self._on_error)
            self.connector_thread.start()
        finally:
            self.mutex.unlock()

    def _on_connected(self, sock: socket.socket):
        self.sock = sock
        self.listener_thread = WiFiListenerThread(self.sock)
        self.listener_thread.data_received.connect(self.data_received)
        self.listener_thread.start()
        self.connected.emit()

    def _on_error(self, error_msg: str):
        self.error.emit(f"连接失败: {error_msg}")

    def disconnect(self):
        self.mutex.lock()
        try:
            if self.listener_thread:
                self.listener_thread.terminate()
                self.listener_thread.wait()
            if self.sock:
                self.sock.close()
            self.disconnected.emit()
        finally:
            self.mutex.unlock()

    def send_command(self, cmd, terminator="\r\n", encoding="utf-8"):
        try:
            data = f"{cmd}{terminator}".encode(encoding)
            if self.sock:
                self.sock.sendall(data)
        except Exception as e:
            self.error.emit(f"发送错误: {str(e)}")


class WiFiConnectorThread(QThread):
    connected = pyqtSignal(object)
    error = pyqtSignal(str)

    def __init__(self, ip, port):
        super().__init__()
        self.ip = ip
        self.port = port

    def run(self):
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(3)
            sock.connect((self.ip, self.port))
            self.connected.emit(sock)
        except Exception as e:
            self.error.emit(str(e))


class WiFiListenerThread(QThread):
    data_received = pyqtSignal(bytes)

    def __init__(self, sock):
        super().__init__()
        self.sock = sock
        self.running = True

    def run(self):
        while self.running:
            try:
                data = self.sock.recv(4096)
                if data:
                    self.data_received.emit(data)
                else:
                    break
            except Exception as e:
                break

    def terminate(self):
        self.running = False


# ================== 数据解析模块 ==================
class FrameParser:
    HEADER = b'\xaa\x55'
    CRC16 = mkCrcFun(0x18005, rev=True, initCrc=0xFFFF, xorOut=0x0000)

    def __init__(self):
        self.buffer = b''

    def feed_data(self, data):
        self.buffer += data
        frames = []
        while len(self.buffer) >= 64:
            candidate = self.buffer[:64]
            if self._validate_frame(candidate):
                frames.append(self._parse_frame(candidate))
                self.buffer = self.buffer[64:]
            else:
                self.buffer = self.buffer[1:]
        return frames

    def _validate_frame(self, frame):
        if len(frame) < 4:
            return False
        calc_crc = self.CRC16(frame[:-2])
        recv_crc = struct.unpack('<H', frame[-2:])[0]
        return frame.startswith(self.HEADER) and calc_crc == recv_crc

    def _parse_frame(self, frame):
        parts = struct.unpack('<2sH6sI8H2xH', frame)
        return {
            'sn': parts[2].decode('utf-8').rstrip('\x00'),
            'timestamp': parts[3],
            'channels': parts[4:12]
        }


# ================== 数据存储模块 ==================
class DataSaver:
    def __init__(self, sn):
        self.sn = sn
        self.buffer = []
        self.lock = threading.Lock()

    def add_frame(self, frame):
        with self.lock:
            self.buffer.append(frame)

    def save_to_mat(self):
        with self.lock:
            if not self.buffer:
                return

            data = {
                'sn': self.sn,
                'timestamps': [f['timestamp'] for f in self.buffer],
                'channels': np.array([f['channels'] for f in self.buffer], dtype=np.uint16)
            }
            savemat(f"{self.sn}_{int(time.time())}.mat", data)


# ================== 主界面 ==================
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.plot_timer = QTimer()
        self.channel_buffers = [[] for _ in range(8)]
        self.init_ui()
        self.communicator = None
        self.parser = FrameParser()
        self.data_saver = None

        self.setWindowTitle("sEMG数据采集系统")
        self.setGeometry(100, 100, 1000, 800)

    def init_ui(self):
        main_widget = QWidget()
        layout = QVBoxLayout()

        # 通信方式选择
        self.comm_type = QComboBox()
        self.comm_type.addItems(["串口", "WiFi"])
        self.comm_stack = QStackedWidget()

        # 串口参数面板
        serial_widget = QWidget()
        serial_layout = QHBoxLayout()
        self.serial_port = QComboBox()
        self.serial_port.addItems([p.portName() for p in QSerialPortInfo.availablePorts()])
        self.baudrate = QComboBox()
        self.baudrate.addItems(["9600", "115200", "230400", "921600", "3686400"])
        serial_layout.addWidget(QLabel("端口:"))
        serial_layout.addWidget(self.serial_port)
        serial_layout.addWidget(QLabel("波特率:"))
        serial_layout.addWidget(self.baudrate)
        serial_widget.setLayout(serial_layout)
        self.comm_stack.addWidget(serial_widget)

        # WiFi参数面板
        wifi_widget = QWidget()
        wifi_layout = QHBoxLayout()
        self.ip_input = QLineEdit("192.168.1.100")
        self.port_input = QLineEdit("5000")
        wifi_layout.addWidget(QLabel("IP:"))
        wifi_layout.addWidget(self.ip_input)
        wifi_layout.addWidget(QLabel("端口:"))
        wifi_layout.addWidget(self.port_input)
        wifi_widget.setLayout(wifi_layout)
        self.comm_stack.addWidget(wifi_widget)

        # 连接控制
        self.conn_btn = QPushButton("连接")
        self.sn_label = QLabel("设备SN: 未连接")

        # 实时绘图区
        plot_group = QGroupBox("实时波形")
        plot_layout = QVBoxLayout()
        self.figure = Figure()
        self.canvas = FigureCanvas(self.figure)
        self.ax = self.figure.add_subplot(111)
        self.lines = [self.ax.plot([], [], lw=0.5)[0] for _ in range(8)]
        self.ax.set_xlim(0, 100)
        self.ax.set_ylim(0, 4096)
        plot_layout.addWidget(self.canvas)
        plot_group.setLayout(plot_layout)

        # 数据采集控制
        data_group = QGroupBox("数据采集")
        data_layout = QHBoxLayout()
        self.duration_input = QLineEdit("1000")
        self.start_btn = QPushButton("开始采集")
        data_layout.addWidget(QLabel("时长(ms):"))
        data_layout.addWidget(self.duration_input)
        data_layout.addWidget(self.start_btn)
        data_group.setLayout(data_layout)

        # 组合布局
        layout.addWidget(self.comm_type)
        layout.addWidget(self.comm_stack)
        layout.addWidget(self.conn_btn)
        layout.addWidget(self.sn_label)
        layout.addWidget(plot_group)
        layout.addWidget(data_group)

        main_widget.setLayout(layout)
        self.setCentralWidget(main_widget)

        # 信号连接
        self.comm_type.currentIndexChanged.connect(self.comm_stack.setCurrentIndex)
        self.conn_btn.clicked.connect(self.toggle_connection)
        self.start_btn.clicked.connect(self.start_capture)
        self.plot_timer.timeout.connect(self.update_plot)
        self.plot_timer.start(50)

    @pyqtSlot()
    def toggle_connection(self):
        if self.conn_btn.text() == "连接":
            self.connect_device()
        else:
            self.disconnect_device()

    def connect_device(self):
        try:
            if self.comm_type.currentText() == "串口":
                port = self.serial_port.currentText()
                baud = int(self.baudrate.currentText())
                self.communicator = SerialCommunicator(port, baud)
            else:
                ip = self.ip_input.text()
                port = int(self.port_input.text())
                self.communicator = WiFiCommunicator(ip, port)

            self.communicator.connected.connect(self.on_connected)
            self.communicator.disconnected.connect(self.on_disconnected)
            self.communicator.data_received.connect(self.handle_data)
            self.communicator.error.connect(lambda msg: self.statusBar().showMessage(msg, 5000))
            self.communicator.connect()

        except Exception as e:
            self.statusBar().showMessage(f"连接错误: {str(e)}", 5000)

    def disconnect_device(self):
        if self.communicator:
            self.communicator.disconnect()

    def send_command(self, cmd):
        if self.communicator:
            self.communicator.send_command(cmd)

    def on_connected(self):
        self.conn_btn.setText("断开连接")
        self.statusBar().showMessage("连接成功!", 2000)
        self.send_command("GET_SN")

    def on_disconnected(self):
        self.conn_btn.setText("连接")
        self.sn_label.setText("设备SN: 未连接")
        self.statusBar().showMessage("已断开连接", 2000)

    @pyqtSlot(bytes)
    def handle_data(self, data):
        frames = self.parser.feed_data(data)
        for frame in frames:
            if 'sn' in frame:
                self.sn_label.setText(f"设备SN: {frame['sn']}")
                if self.data_saver:
                    self.data_saver.sn = frame['sn']
            for i in range(8):
                self.channel_buffers[i].append(frame['channels'][i])
                if len(self.channel_buffers[i]) > 100:
                    self.channel_buffers[i] = self.channel_buffers[i][-100:]
            if self.data_saver:
                self.data_saver.add_frame(frame)

    def update_plot(self):
        xdata = np.arange(100)
        for i, line in enumerate(self.lines):
            ydata = self.channel_buffers[i][-100:] or [0] * 100
            line.set_data(xdata[-len(ydata):], ydata)
        self.ax.relim()
        self.ax.autoscale_view()
        self.canvas.draw()

    def start_capture(self):
        try:
            duration = int(self.duration_input.text())
            self.send_command("START_DATA")
            self.data_saver = DataSaver(self.sn_label.text().split(": ")[-1])
            QTimer.singleShot(duration, self.stop_capture)
            self.statusBar().showMessage("数据采集中...")
        except ValueError:
            self.statusBar().showMessage("无效的时长参数", 3000)

    def stop_capture(self):
        self.send_command("STOP_DATA")
        if self.data_saver:
            self.data_saver.save_to_mat()
            self.data_saver = None
            self.statusBar().showMessage("数据保存成功!", 5000)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())